{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from lightrag.core import Component, Generator\n",
    "from lightrag.components.model_client import OpenAIClient\n",
    "from lightrag.components.model_client import GroqAPIClient\n",
    "from lightrag.utils import setup_env # make sure you have a .env file with OPENAI_API_KEY and GROQ_API_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_doc = r\"\"\"<SYS> You are a doctor </SYS> User: {{input_str}}\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's turn on the library log to help with debugging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lightrag.utils import get_logger\n",
    "get_logger()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Toy example\n",
    "\n",
    "class DocQA(Component):\n",
    "    def __init__(self):\n",
    "        super(DocQA, self).__init__()\n",
    "        self.doc = Generator(\n",
    "            template=template_doc,\n",
    "            model_client=OpenAIClient(),\n",
    "            model_kwargs={\"model\": \"gpt-3.5-turbo\"},\n",
    "        )\n",
    "\n",
    "    def call(self, query: str) -> str:\n",
    "        return self.doc(prompt_kwargs={\"input_str\": query}).data\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-17 09:42:27 - INFO - [prompt_builder.py:82:__init__] - Prompt has variables: []\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DocQA(\n",
       "  (doc): Generator(\n",
       "    model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
       "    (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}})\n",
       "    (model_client): OpenAIClient()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc = DocQA()\n",
    "doc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'type': 'DocQA', 'data': {'_components': {'_ordered_dict': True, 'data': [('doc', {'type': 'Generator', 'data': {'_components': {'_ordered_dict': True, 'data': [('prompt', {'type': 'Prompt', 'data': {'_components': {'_ordered_dict': True, 'data': []}, '_parameters': {'_ordered_dict': True, 'data': []}, 'training': False, 'template': '<SYS> You are a doctor </SYS> User: {{input_str}}', 'prompt_variables': [], 'preset_prompt_kwargs': {}}}), ('model_client', {'type': 'OpenAIClient', 'data': {'_components': {'_ordered_dict': True, 'data': []}, '_parameters': {'_ordered_dict': True, 'data': []}, 'training': False, '_api_key': None}})]}, '_parameters': {'_ordered_dict': True, 'data': []}, 'training': False, 'template': '<SYS> You are a doctor </SYS> User: {{input_str}}', 'preset_prompt_kwargs': {}, 'model_kwargs': {'model': 'gpt-3.5-turbo'}, 'output_processors': None, '_trainable_params': []}})]}, '_parameters': {'_ordered_dict': True, 'data': []}, 'training': False}}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'_components': OrderedDict([('doc',\n",
       "               Generator(\n",
       "                 model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
       "                 (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}})\n",
       "                 (model_client): OpenAIClient()\n",
       "               ))]),\n",
       " '_parameters': OrderedDict(),\n",
       " 'training': False}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# states\n",
    "states = doc.to_dict()\n",
    "print(states)\n",
    "doc.__dict__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_components': OrderedDict([('doc',\n",
       "               Generator(\n",
       "                 model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
       "                 (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}})\n",
       "                 (model_client): OpenAIClient()\n",
       "               ))]),\n",
       " '_parameters': OrderedDict(),\n",
       " 'training': False}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# restore the states\n",
    "doc2 = DocQA.from_dict(states)\n",
    "# print(doc2.call(\"What is the capital of France?\"))\n",
    "doc2.__dict__\n",
    "# doc2.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'type': 'DocQA',\n",
       " 'data': {'_components': {'_ordered_dict': True,\n",
       "   'data': [('doc',\n",
       "     {'type': 'Generator',\n",
       "      'data': {'_components': {'_ordered_dict': True,\n",
       "        'data': [('prompt',\n",
       "          {'type': 'Prompt',\n",
       "           'data': {'_components': {'_ordered_dict': True, 'data': []},\n",
       "            '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "            'training': False,\n",
       "            'template': '<SYS> You are a doctor </SYS> User: {{input_str}}',\n",
       "            'prompt_variables': [],\n",
       "            'preset_prompt_kwargs': {}}}),\n",
       "         ('model_client',\n",
       "          {'type': 'OpenAIClient',\n",
       "           'data': {'_components': {'_ordered_dict': True, 'data': []},\n",
       "            '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "            'training': False,\n",
       "            '_api_key': None}})]},\n",
       "       '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "       'training': False,\n",
       "       'template': '<SYS> You are a doctor </SYS> User: {{input_str}}',\n",
       "       'preset_prompt_kwargs': {},\n",
       "       'model_kwargs': {'model': 'gpt-3.5-turbo'},\n",
       "       'output_processors': None,\n",
       "       '_trainable_params': []}})]},\n",
       "  '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "  'training': False}}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc2.to_dict() == doc.to_dict()\n",
    "doc2.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pickle to states\n",
    "import pickle\n",
    "# from collections import OrderedDict\n",
    "# from openai import OpenAI # cant pickle this\n",
    "\n",
    "# class DummpyDocQA():\n",
    "#     a = OrderedDict()\n",
    "#     def __init__(self):\n",
    "#         self.dummpy = 1\n",
    "#         self.client = OpenAI()\n",
    "\n",
    "# doc_dummy = DummpyDocQA()\n",
    "with open(\"doc.pkl\", \"wb\") as f:\n",
    "    pickle.dump(doc.to_dict(), f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the serialized states to a file\n",
    "from lightrag.utils.file_io import save_pickle, save_json\n",
    "states = doc.to_dict()\n",
    "# save_json(states, \"doc.json\")\n",
    "save_pickle(states, \"doc.pkl\")\n",
    "\n",
    "# load the serialized states from a file\n",
    "from lightrag.utils.file_io import load_pickle, load_json\n",
    "states_loaded = load_pickle(\"doc.pkl\")\n",
    "# states_loaded = load_json(\"doc.json\")\n",
    "\n",
    "states_loaded == states\n",
    "\n",
    "doc3 = DocQA.from_dict(states_loaded)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-14 17:42:48 - INFO - [generator.py:199:call] - prompt_kwargs: {'input_str': 'What is the capital of France?'}\n",
      "2024-06-14 17:42:48 - INFO - [generator.py:200:call] - model_kwargs: {}\n",
      "2024-06-14 17:42:48 - WARNING - [prompt_builder.py:120:compose_prompt_kwargs] - Key input_str does not exist in the prompt_kwargs.\n",
      "2024-06-14 17:42:48 - INFO - [openai_client.py:139:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a doctor </SYS> User: What is the capital of France?'}]}\n",
      "2024-06-14 17:42:48 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-14 17:42:48 - INFO - [generator.py:208:call] - output: GeneratorOutput(data='The capital of France is Paris.', error=None, usage=None, raw_response='The capital of France is Paris.')\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The capital of France is Paris.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc3\n",
    "doc3.call(\"What is the capital of France?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-14 17:12:51 - INFO - [generator.py:199:call] - prompt_kwargs: {'input_str': 'What is the best treatment for headache?'}\n",
      "2024-06-14 17:12:51 - INFO - [generator.py:200:call] - model_kwargs: {}\n",
      "2024-06-14 17:12:51 - INFO - [openai_client.py:140:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a doctor </SYS> User: What is the best treatment for headache?'}]}\n",
      "2024-06-14 17:12:54 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-14 17:12:54 - INFO - [generator.py:208:call] - output: GeneratorOutput(data='As a doctor, the best treatment for a headache depends on the cause of the headache. In general, some common treatments for headaches include:\\n\\n1. Over-the-counter pain relievers such as acetaminophen, ibuprofen, or aspirin\\n2. Rest and relaxation in a quiet, dark room\\n3. Hydration\\n4. Applying a cold or warm compress to the forehead or neck\\n5. Avoiding triggers such as stress, lack of sleep, or certain foods\\n6. Practicing relaxation techniques such as deep breathing, meditation, or yoga\\n7. Prescription medications for chronic or severe headaches\\n\\nIt is important to consult with a healthcare provider for a proper diagnosis and treatment plan for your specific type of headache.', error=None, usage=None, raw_response='As a doctor, the best treatment for a headache depends on the cause of the headache. In general, some common treatments for headaches include:\\n\\n1. Over-the-counter pain relievers such as acetaminophen, ibuprofen, or aspirin\\n2. Rest and relaxation in a quiet, dark room\\n3. Hydration\\n4. Applying a cold or warm compress to the forehead or neck\\n5. Avoiding triggers such as stress, lack of sleep, or certain foods\\n6. Practicing relaxation techniques such as deep breathing, meditation, or yoga\\n7. Prescription medications for chronic or severe headaches\\n\\nIt is important to consult with a healthcare provider for a proper diagnosis and treatment plan for your specific type of headache.')\n",
      "As a doctor, the best treatment for a headache depends on the cause of the headache. In general, some common treatments for headaches include:\n",
      "\n",
      "1. Over-the-counter pain relievers such as acetaminophen, ibuprofen, or aspirin\n",
      "2. Rest and relaxation in a quiet, dark room\n",
      "3. Hydration\n",
      "4. Applying a cold or warm compress to the forehead or neck\n",
      "5. Avoiding triggers such as stress, lack of sleep, or certain foods\n",
      "6. Practicing relaxation techniques such as deep breathing, meditation, or yoga\n",
      "7. Prescription medications for chronic or severe headaches\n",
      "\n",
      "It is important to consult with a healthcare provider for a proper diagnosis and treatment plan for your specific type of headache.\n"
     ]
    }
   ],
   "source": [
    "print(doc(\"What is the best treatment for headache?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-14 17:12:54 - INFO - [generator.py:199:call] - prompt_kwargs: {'input_str': 'What is the best treatment for headache?'}\n",
      "2024-06-14 17:12:54 - INFO - [generator.py:200:call] - model_kwargs: {}\n",
      "2024-06-14 17:12:54 - INFO - [openai_client.py:140:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a doctor </SYS> User: What is the best treatment for headache?'}]}\n",
      "2024-06-14 17:12:56 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-14 17:12:56 - INFO - [generator.py:208:call] - output: GeneratorOutput(data='As a doctor, the best treatment for a headache can depend on the cause of the headache. If the headache is mild and infrequent, over-the-counter pain relievers such as ibuprofen or acetaminophen can help. Additionally, getting enough rest, staying hydrated, and practicing relaxation techniques such as deep breathing exercises or meditation can also provide relief. If the headache is severe, persistent, or accompanied by other concerning symptoms, it is important to consult a healthcare professional for a proper diagnosis and individualized treatment plan.', error=None, usage=None, raw_response='As a doctor, the best treatment for a headache can depend on the cause of the headache. If the headache is mild and infrequent, over-the-counter pain relievers such as ibuprofen or acetaminophen can help. Additionally, getting enough rest, staying hydrated, and practicing relaxation techniques such as deep breathing exercises or meditation can also provide relief. If the headache is severe, persistent, or accompanied by other concerning symptoms, it is important to consult a healthcare professional for a proper diagnosis and individualized treatment plan.')\n",
      "As a doctor, the best treatment for a headache can depend on the cause of the headache. If the headache is mild and infrequent, over-the-counter pain relievers such as ibuprofen or acetaminophen can help. Additionally, getting enough rest, staying hydrated, and practicing relaxation techniques such as deep breathing exercises or meditation can also provide relief. If the headache is severe, persistent, or accompanied by other concerning symptoms, it is important to consult a healthcare professional for a proper diagnosis and individualized treatment plan.\n"
     ]
    }
   ],
   "source": [
    "print(doc2(\"What is the best treatment for headache?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('', DocQA(\n",
      "  (doc): Generator(\n",
      "    model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
      "    (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
      "    (model_client): OpenAIClient()\n",
      "  )\n",
      "))\n",
      "('doc', Generator(\n",
      "  model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
      "  (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
      "  (model_client): OpenAIClient()\n",
      "))\n",
      "('doc.prompt', Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str']))\n",
      "('doc.model_client', OpenAIClient())\n"
     ]
    }
   ],
   "source": [
    "# list other subcomponents\n",
    "\n",
    "for subcomponent in doc.named_components():\n",
    "    print(subcomponent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's add a parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lightrag.core.parameter import Parameter\n",
    "\n",
    "doc.register_parameter(\"demo\", param=Parameter(data=\"demo\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('demo', Parameter: demo)\n"
     ]
    }
   ],
   "source": [
    "# list all parameters\n",
    "for param in doc.named_parameters():\n",
    "    print(param)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'type': 'DocQA',\n",
       " 'data': {'_components': {'_ordered_dict': True,\n",
       "   'data': [('doc',\n",
       "     {'type': 'Generator',\n",
       "      'data': {'_components': {'_ordered_dict': True,\n",
       "        'data': [('prompt',\n",
       "          {'type': 'Prompt',\n",
       "           'data': {'_components': {'_ordered_dict': True, 'data': []},\n",
       "            '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "            'training': False,\n",
       "            '_template_string': '<SYS> You are a doctor </SYS> User: {{input_str}}',\n",
       "            'template': <Template memory:15f114f50>,\n",
       "            'prompt_variables': ['input_str'],\n",
       "            'preset_prompt_kwargs': {}}}),\n",
       "         ('model_client',\n",
       "          {'type': 'OpenAIClient',\n",
       "           'data': {'_components': {'_ordered_dict': True, 'data': []},\n",
       "            '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "            'training': False,\n",
       "            'sync_client': <openai.OpenAI at 0x15cfc1cd0>,\n",
       "            'async_client': None,\n",
       "            '_api_key': None}})]},\n",
       "       '_parameters': {'_ordered_dict': True, 'data': []},\n",
       "       'training': False,\n",
       "       'template': '<SYS> You are a doctor </SYS> User: {{input_str}}',\n",
       "       'preset_prompt_kwargs': {},\n",
       "       'model_kwargs': {'model': 'gpt-3.5-turbo'},\n",
       "       'output_processors': None,\n",
       "       '_trainable_params': []}})]},\n",
       "  '_parameters': {'_ordered_dict': True,\n",
       "   'data': [('demo', {'data': 'demo', 'requires_opt': True})]},\n",
       "  'training': False}}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lightrag.utils.file_io import save_json\n",
    "\n",
    "save_json(doc.to_dict(), \"doc.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('demo', Parameter: demo)])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-14 17:12:56 - INFO - [generator.py:199:call] - prompt_kwargs: {'input_str': 'What is the best treatment for a cold?'}\n",
      "2024-06-14 17:12:56 - INFO - [generator.py:200:call] - model_kwargs: {}\n",
      "2024-06-14 17:12:56 - INFO - [openai_client.py:140:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a doctor </SYS> User: What is the best treatment for a cold?'}]}\n",
      "2024-06-14 17:12:57 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-14 17:12:57 - INFO - [generator.py:208:call] - output: GeneratorOutput(data=\"As a doctor, I recommend a combination of rest, staying hydrated, over-the-counter cold medications (such as decongestants or pain relievers), throat lozenges, and steam inhalation. It's also important to eat a balanced diet, get plenty of rest, and avoid close contact with others to prevent spreading the cold. If symptoms persist or worsen, it's best to consult with a healthcare provider for further evaluation and treatment.\", error=None, usage=None, raw_response=\"As a doctor, I recommend a combination of rest, staying hydrated, over-the-counter cold medications (such as decongestants or pain relievers), throat lozenges, and steam inhalation. It's also important to eat a balanced diet, get plenty of rest, and avoid close contact with others to prevent spreading the cold. If symptoms persist or worsen, it's best to consult with a healthcare provider for further evaluation and treatment.\")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"As a doctor, I recommend a combination of rest, staying hydrated, over-the-counter cold medications (such as decongestants or pain relievers), throat lozenges, and steam inhalation. It's also important to eat a balanced diet, get plenty of rest, and avoid close contact with others to prevent spreading the cold. If symptoms persist or worsen, it's best to consult with a healthcare provider for further evaluation and treatment.\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "doc.call(\"What is the best treatment for a cold?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "<class 'lightrag.core.component.FunComponent'>\n"
     ]
    }
   ],
   "source": [
    "from lightrag.core.component import FunComponent\n",
    "\n",
    "def add_one(x):\n",
    "    return x + 1\n",
    "\n",
    "fun_component = FunComponent(add_one)\n",
    "print(fun_component(1))  \n",
    "print(type(fun_component))  \n",
    "\n",
    "# output:\n",
    "# 2\n",
    "# <class 'core.component.FunComponent'>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "<class 'lightrag.core.component.AddOneComponent'>\n"
     ]
    }
   ],
   "source": [
    "from lightrag.core.component import fun_to_component \n",
    "\n",
    "fun_component = fun_to_component(add_one)\n",
    "print(fun_component(1))\n",
    "print(type(fun_component))\n",
    "\n",
    "# output:\n",
    "# 2\n",
    "# <class 'lightrag.core.component.AddOneComponent'>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "<class 'lightrag.core.component.AddOneComponent'>\n"
     ]
    }
   ],
   "source": [
    "# use it as a decorator\n",
    "@fun_to_component\n",
    "def add_one(x):\n",
    "    return x + 1\n",
    "\n",
    "print(add_one(1))\n",
    "print(type(add_one))\n",
    "\n",
    "# output:\n",
    "# 2\n",
    "# <class 'lightrag.core.component.AddOneComponent'>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-14 17:12:57 - INFO - [generator.py:199:call] - prompt_kwargs: {'input_str': 'What is the best treatment for headache?Please be concise and only list the top treatments.'}\n",
      "2024-06-14 17:12:57 - INFO - [generator.py:200:call] - model_kwargs: {}\n",
      "2024-06-14 17:12:57 - INFO - [openai_client.py:140:call] - api_kwargs: {'model': 'gpt-3.5-turbo', 'messages': [{'role': 'system', 'content': '<SYS> You are a doctor </SYS> User: What is the best treatment for headache?Please be concise and only list the top treatments.'}]}\n",
      "2024-06-14 17:12:58 - INFO - [_client.py:1026:_send_single_request] - HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "2024-06-14 17:12:58 - INFO - [generator.py:208:call] - output: GeneratorOutput(data='1. Over-the-counter pain relievers such as ibuprofen or acetaminophen\\n2. Stay hydrated and rest\\n3. Apply a cold compress to the forehead\\n4. Practice relaxation techniques such as deep breathing or meditation', error=None, usage=None, raw_response='1. Over-the-counter pain relievers such as ibuprofen or acetaminophen\\n2. Stay hydrated and rest\\n3. Apply a cold compress to the forehead\\n4. Practice relaxation techniques such as deep breathing or meditation')\n",
      "1. Over-the-counter pain relievers such as ibuprofen or acetaminophen\n",
      "2. Stay hydrated and rest\n",
      "3. Apply a cold compress to the forehead\n",
      "4. Practice relaxation techniques such as deep breathing or meditation\n"
     ]
    }
   ],
   "source": [
    "from lightrag.core.component import Sequential\n",
    "\n",
    "@fun_to_component\n",
    "def enhance_query(query:str) -> str:\n",
    "    return query + \"Please be concise and only list the top treatments.\"\n",
    "\n",
    "seq = Sequential(enhance_query, doc)\n",
    "\n",
    "query = \"What is the best treatment for headache?\"\n",
    "print(seq(query))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sequential with just a function, will raise error\n",
    "# def enhance_query(query:str) -> str:\n",
    "#     return query + \"Please be concise and only list the top treatments.\"\n",
    "# seq2 = Sequential(enhance_query, doc)\n",
    "# print(seq2(query))\n",
    "# print(seq2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): EnhanceQueryComponent()\n",
       "  (1): DocQA(\n",
       "    (doc): Generator(\n",
       "      model_kwargs={'model': 'gpt-3.5-turbo'}, \n",
       "      (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])\n",
       "      (model_client): OpenAIClient()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TODO: LLM for single choices"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
